package com.breeze.components.core.util;

import cn.hutool.core.collection.CollUtil;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.function.BiConsumer;
import java.util.function.Function;
import java.util.stream.Collectors;
import lombok.AccessLevel;
import lombok.NoArgsConstructor;
import lombok.extern.slf4j.Slf4j;

/**
 * @author breeze
 */
@Slf4j
@NoArgsConstructor(access = AccessLevel.PRIVATE)
public class TreeUtil extends cn.hutool.core.lang.tree.TreeUtil {

  public static <T, R> List<T> buildTree(
      List<T> nodeList,
      Function<T, R> getCurrentId,
      Function<T, R> getParentId,
      R rootValue,
      BiConsumer<T, List<T>> setChildren) {
    if (CollUtil.isEmpty(nodeList)) {
      return Collections.emptyList();
    }
    Map<R, List<T>> parentIdMap = nodeList.stream().collect(Collectors.groupingBy(getParentId));
    List<T> root = parentIdMap.get(rootValue);
    deepSetChildren(root, getCurrentId, parentIdMap, setChildren);
    return root;
  }

  private static <T, R> void deepSetChildren(
      List<T> nodeList,
      Function<T, R> getCurrentId,
      Map<R, List<T>> parentIdMap,
      BiConsumer<T, List<T>> setChildren) {
    if (CollUtil.isEmpty(nodeList)) {
      return;
    }
    nodeList.forEach(
        s -> {
          List<T> innerTs = parentIdMap.get(getCurrentId.apply(s));
          deepSetChildren(innerTs, getCurrentId, parentIdMap, setChildren);
          setChildren.accept(s, innerTs);
        });
  }
}
